import pandas as pd
import numpy as np
import pulp
import math
import os
from typing import Dict

class ProductionScheduler:
    def __init__(self, data_file='data/data.xlsx'):
        # 获取当前文件所在目录（utils文件夹）
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 获取app.py所在目录（上级目录）
        app_dir = os.path.dirname(current_dir)
        # 构建正确的文件路径
        data_path = os.path.join(app_dir, data_file)
        print(f"正在读取数据文件: {data_path}")
        
        # 检查文件是否存在
        if not os.path.exists(data_path):
            raise FileNotFoundError(f"数据文件不存在: {data_path}")
        
        self.df = pd.read_excel(data_path)
        self.pivot_df = self.df.pivot(index='工龄', columns='生产线编号', values='预测净合格数')
        self.num_lines = len(self.pivot_df.columns)
        self.total_workers_needed = math.ceil(self.num_lines * 7 * 3 / 5)
        self.ages = sorted(self.pivot_df.index.unique())
        
        print(f"数据加载成功: {len(self.df)} 行, {self.num_lines} 条生产线")
        print(f"需要工人总数: {self.total_workers_needed}")


    def schedule_production(self, age_to_count: Dict[int, int]):
        """
        执行生产排班
        """
        # 验证总数
        total_count = sum(age_to_count.values())
        if total_count != self.total_workers_needed:
            raise ValueError(f"工人总数必须为 {self.total_workers_needed}，当前为 {total_count}")
        
        # 检查数据完整性 - 先转换键类型为整数
        age_to_count_int = {}
        for key, value in age_to_count.items():
            try:
                age_to_count_int[int(key)] = int(value)
            except (ValueError, TypeError):
                raise ValueError(f"无效的工龄或人数: {key}={value}")
        
        data_ages = set(self.pivot_df.index)
        input_ages = set(age_to_count_int.keys())
        
        # 如果有缺失的工龄，自动补0
        missing_ages = data_ages - input_ages
        if missing_ages:
            print(f"警告: 输入缺少工龄 {missing_ages}，自动设置为0")
            for age in missing_ages:
                age_to_count_int[age] = 0
        
        # 使用转换后的字典 - 这里是最关键的修改！
        repeat_counts = [age_to_count_int[age] for age in self.pivot_df.index]
        
        # 复制数据
        data_repeated = self.pivot_df.loc[self.pivot_df.index.repeat(repeat_counts)]
        
        # 创建Net_Qua矩阵
        Net_Qua = data_repeated.values
        num_workers = total_count
        
        # 定义问题模型
        problem = pulp.LpProblem("Maximize Net Qualified Quantity", pulp.LpMaximize)
        
        # 定义变量
        num_shifts = 3
        num_days = 7
        x = pulp.LpVariable.dicts('x', 
                                 ((i, j, k, t) for i in range(num_workers) 
                                  for j in range(self.num_lines) 
                                  for k in range(num_shifts) 
                                  for t in range(num_days)), 
                                 cat='Binary')
        
        # 辅助变量
        y = pulp.LpVariable.dicts('y', 
                                 ((i, j) for i in range(num_workers) for j in range(7)), 
                                 cat='Binary')
        
        # 目标函数
        problem += pulp.lpSum([x[i, j, k, t] * Net_Qua[i][j] 
                             for i in range(num_workers) 
                             for j in range(self.num_lines) 
                             for k in range(num_shifts) 
                             for t in range(num_days)])
        
        # 约束条件
        # 每个班次只能由一个工人操作
        for j in range(self.num_lines):
            for k in range(num_shifts):
                for t in range(num_days):
                    problem += pulp.lpSum([x[i, j, k, t] for i in range(num_workers)]) == 1
        
        # 每个工人每天只能工作一个班次
        for i in range(num_workers):
            for t in range(num_days):
                problem += pulp.lpSum([x[i, j, k, t] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1
        
        # 工人最多工作5天
        for i in range(num_workers):
            problem += pulp.lpSum([x[i, j, k, t] for j in range(self.num_lines) for k in range(num_shifts) for t in range(7)]) <= 5
        
        # 工人如果前一天是晚班，则第二天不能安排早班
        for i in range(num_workers):
            for t in range(num_days):
                problem += pulp.lpSum([x[i, j, 2, t] for j in range(self.num_lines)]) + \
                          pulp.lpSum([x[i, j, 0, (t+1)%7] for j in range(self.num_lines)]) <= 1
        
        # 休息安排约束
        for n in range(num_workers):
            # 约束1: 1-5天工作，6-7天休息
            for i in range(0, 5):
                problem += pulp.lpSum([x[n, j, k, i] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n, 0]
            problem += pulp.lpSum([x[n, j, k, 5] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n, 0]
            problem += pulp.lpSum([x[n, j, k, 6] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n, 0]
            
        # 约束2: 2-6天工作，1、7天休息
        for n in range(42):
            for i in range(1, 6):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,1]
            problem += pulp.lpSum([x[n, j, k, 0 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,1]
            problem += pulp.lpSum([x[n, j, k, 6 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,1]
            
            
            
            # 约束3: 3-7天工作，1-2天休息
            for i in range(2, 7):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,2]
            problem += pulp.lpSum([x[n, j, k, 0 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,2]
            problem += pulp.lpSum([x[n, j, k, 1 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,2]
            
            # 约束4: 1，,4-7天工作，2-3天休息
            for i in range(3, 7):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,3]
            problem += pulp.lpSum([x[n, j, k, 0 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,3]
            problem += pulp.lpSum([x[n, j, k, 1 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,3]
            problem += pulp.lpSum([x[n, j, k, 2 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,3]
            
            # 约束5: 1，,4-7天工作，2-3天休息
            for i in range(4, 7):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,4]
            problem += pulp.lpSum([x[n, j, k, 0 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,4]
            problem += pulp.lpSum([x[n, j, k, 1 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,4]
            problem += pulp.lpSum([x[n, j, k, 2 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,4]
            problem += pulp.lpSum([x[n, j, k, 3 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,4]
            
            # 约束6:
            for i in range(0, 3):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,5]
            problem += pulp.lpSum([x[n, j, k, 5 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,5]
            problem += pulp.lpSum([x[n, j, k, 6 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,5]
            problem += pulp.lpSum([x[n, j, k, 3 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,5]
            problem += pulp.lpSum([x[n, j, k, 4 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,5]
            
            # 约束7:
            for i in range(0, 4):
                problem+= pulp.lpSum([x[n, j, k, i ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,6]
            
            problem += pulp.lpSum([x[n, j, k, 6 ] for j in range(self.num_lines) for k in range(num_shifts)]) >= y[n,6]
            problem += pulp.lpSum([x[n, j, k, 4 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,6]
            problem += pulp.lpSum([x[n, j, k, 5 ] for j in range(self.num_lines) for k in range(num_shifts)]) <= 1 - y[n,6]

            
            # 确保只有一种安排被选择
            problem += pulp.lpSum(y[n, j] for j in range(7)) == 1
        
        # 班次尽量均衡
        for i in range(num_workers):
            for shift in range(3):
                problem += pulp.lpSum([x[i, j, shift, t] for j in range(self.num_lines) for t in range(7)]) <= 2
        
        # 求解模型
        problem.solve()
        
        # 生成结果
        result1_df, result2_df = self._generate_results(x, num_workers, self.num_lines, num_shifts, num_days)
        
        return result1_df, result2_df
    
    def _generate_results(self, x, num_workers, num_lines, num_shifts, num_days):
        """生成两种格式的结果"""
        # 生成详细排班结果
        detailed_data = []
        for i in range(num_workers):
            for j in range(num_lines):
                for k in range(num_shifts):
                    for t in range(num_days):
                        if pulp.value(x[i, j, k, t]) > 0:
                            detailed_data.append({
                                '日期': t + 1,
                                '班次': k + 1,
                                '操作人员': i + 1
                            })
        
        detailed_df = pd.DataFrame(detailed_data)
        
        # 生成result4-1格式
        result1_df = self._create_result1_format(detailed_df, num_workers, num_days)
        
        # 生成result4-2格式
        result2_df = self._create_result2_format(detailed_df)
        
        return result1_df, result2_df
    
    def _create_result1_format(self, detailed_df, num_workers, num_days):
        """创建result4-1格式的数据"""
        # 创建空的数据框
        columns = ['日期'] + [f'B{str(i+1).zfill(3)}' for i in range(num_workers)]
        result1_data = []
        
        for day in range(1, num_days + 1):
            day_data = {'日期': day}
            
            # 获取当天的排班
            day_schedule = detailed_df[detailed_df['日期'] == day]
            
            # 初始化所有工人为休息
            for worker in range(1, num_workers + 1):
                day_data[f'B{str(worker).zfill(3)}'] = '休'
            
            # 填充工作班次
            for _, row in day_schedule.iterrows():
                worker_id = row['操作人员']
                shift = row['班次']
                shift_name = {1: '早', 2: '中', 3: '晚'}[shift]
                day_data[f'B{str(worker_id).zfill(3)}'] = shift_name
            
            result1_data.append(day_data)
        
        return pd.DataFrame(result1_data)
    
    def _create_result2_format(self, detailed_df):
        """创建result4-2格式的数据"""
        # 按日期和班次分组
        grouped = detailed_df.groupby(['日期', '班次'])['操作人员'].apply(list).reset_index()
        
        # 将列表转换为多个列
        max_operators = grouped['操作人员'].apply(len).max()
        operator_cols = [f'M{301+i}' for i in range(max_operators)]
        
        operator_df = grouped['操作人员'].apply(pd.Series)
        operator_df.columns = operator_cols
        
        # 合并数据
        result2_df = pd.concat([grouped[['日期', '班次']], operator_df], axis=1)
        
        # 转换班次名称和工人ID格式
        result2_df['班次'] = result2_df['班次'].replace({1: '早', 2: '中', 3: '晚'})
        
        for col in operator_cols:
            result2_df[col] = result2_df[col].apply(
                lambda x: f'B{str(int(x)).zfill(3)}' if pd.notna(x) else None
            )
        
        return result2_df

# 单例实例
scheduler = ProductionScheduler()

def update_schedule(age_distribution):
    """
    更新排班
    
    Args:
        age_distribution: 工龄分布字典，如 {1: 8, 2: 8, 3: 4, 4: 8, 5: 9, 6: 5}
    
    Returns:
        bool: 是否成功
    """
    try:
        print(f"开始排班，工龄分布: {age_distribution}")
        
        result1_df, result2_df = scheduler.schedule_production(age_distribution)
        
        # 获取当前文件所在目录（utils文件夹）
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 获取app.py所在目录（上级目录）
        app_dir = os.path.dirname(current_dir)
        # 构建data目录路径
        data_dir = os.path.join(app_dir, 'data')
        
        # 确保data目录存在
        if not os.path.exists(data_dir):
            os.makedirs(data_dir)
            print(f"创建数据目录: {data_dir}")
        
        # 构建完整的文件路径
        result1_path = os.path.join(data_dir, 'result4-1.xlsx')
        result2_path = os.path.join(data_dir, 'result4-2.xlsx')
        
        # 保存结果
        result1_df.to_excel(result1_path, index=False)
        result2_df.to_excel(result2_path, index=False)
        
        print(f"排班完成，文件已保存到: {result1_path}, {result2_path}")
        return True
        
    except Exception as e:
        print(f"排班更新失败: {e}")
        import traceback
        traceback.print_exc()  # 打印完整的错误堆栈
        return False